import pymysql
from dbutils.pooled_db import PooledDB
from typing import List, Dict, Union, Optional

from dotenv import load_dotenv, find_dotenv
import os
load_dotenv(find_dotenv())

class MySQLDB:
    def __init__(self):
        """
        MySQL数据库工具类
        """
        db_config = {
            'host': os.getenv('db_host'),
            'port': int(os.getenv('db_port', 3306)),
            'user': os.getenv('db_username'),
            'password': os.getenv('db_pwd'),
            'database': os.getenv('db_database'),
            'charset': 'utf8mb4',
            'autocommit': False,
            'cursorclass': pymysql.cursors.DictCursor
        }

        pool_config = {
            'mincached': 1,  # 初始空闲连接数
            'maxcached': 5,  # 最大空闲连接数
            'maxconnections': 10,  # 最大总连接数
            'blocking': True,  # 连接池满时阻塞等待
            'ping': 1  # 每次使用前自动检测连接
        }

        # 初始化连接池
        self.pool = PooledDB(
            creator=pymysql,  # 使用pymysql作为底层驱动
            **db_config,  # 数据库连接参数
            **pool_config  # 连接池参数
        )

    def get_conn(self):
        """获取数据库连接"""
        return self.pool.connection()

    def execute_query(self, sql: str, params: Union[tuple, dict, None] = None) -> List[Dict]:
        """
        执行查询语句
        :param sql: SQL语句，支持参数化查询
        :param params: 参数，可以是元组或字典
        :return: 查询结果列表，字典格式
        """
        with self.get_conn() as conn:
            with conn.cursor(pymysql.cursors.DictCursor) as cursor:
                cursor.execute(sql, params)
                return cursor.fetchall()

    def execute_update(self, sql: str, params: Union[tuple, dict, None] = None) -> int:
        """
        执行增删改操作
        :param sql: SQL语句
        :param params: 参数
        :return: 影响的行数
        """
        with self.get_conn() as conn:
            with conn.cursor() as cursor:
                affected_rows = cursor.execute(sql, params)
                conn.commit()
                return affected_rows

    def execute_many(self, sql: str, params: List[Union[tuple, dict]]) -> int:
        """
        批量执行操作
        :param sql: SQL语句
        :param params: 参数列表
        :return: 影响的行数
        """
        with self.get_conn() as conn:
            with conn.cursor() as cursor:
                affected_rows = cursor.executemany(sql, params)
                conn.commit()
                return affected_rows

    def transaction(self):
        """事务上下文管理器"""
        return self.Transaction(self)

    class Transaction:
        def __init__(self, db):
            self.db = db
            self.conn = None

        def __enter__(self):
            self.conn = self.db.get_conn()
            self.conn.begin()
            return self

        def __exit__(self, exc_type, exc_val, exc_tb):
            if exc_type:
                self.conn.rollback()
            else:
                self.conn.commit()
            self.conn.close()

        def execute(self, sql: str, params: Union[tuple, dict, None] = None) -> int:
            """在事务中执行SQL"""
            with self.conn.cursor() as cursor:
                return cursor.execute(sql, params)

        def executemany(self, sql: str, params: List[Union[tuple, dict]]) -> int:
            """在事务中批量执行SQL"""
            with self.conn.cursor() as cursor:
                return cursor.executemany(sql, params)

    def table_exists(self, table_name: str) -> bool:
        """检查表是否存在"""
        sql = """
        SELECT COUNT(*) 
        FROM information_schema.tables 
        WHERE table_schema = DATABASE() 
        AND table_name = %s
        """
        result = self.execute_query(sql, (table_name,))
        return result[0]['COUNT(*)'] > 0

    def close(self):
        """关闭连接池"""
        self.pool.close()

# 使用示例
if __name__ == '__main__':


    # 初始化数据库连接池
    db = MySQLDB()

    try:
        # 查询示例
        # users = db.execute_query("SELECT * FROM users WHERE age > %s", (18,))
        # print(f"查询结果：{users}")
        #
        # # 插入示例
        # insert_sql = "INSERT INTO users (name, age) VALUES (%s, %s)"
        # new_id = db.execute_update(insert_sql, ("Alice", 25))
        # print(f"插入数据，影响行数：{new_id}")
        #
        # # 批量插入示例
        # batch_data = [("Bob", 30), ("Charlie", 35)]
        # affected = db.execute_many(insert_sql, batch_data)
        # print(f"批量插入影响行数：{affected}")
        #
        # # 事务示例
        # with db.transaction() as tx:
        #     tx.execute("UPDATE accounts SET balance = balance - 100 WHERE user_id = 1")
        #     tx.execute("UPDATE accounts SET balance = balance + 100 WHERE user_id = 2")

        # 表存在检查
        print(f"users表是否存在：{db.table_exists('strategy_stock')}")

    finally:
        db.close()